// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/prestosql/aggregates/min_max_aggregates.h>
#include <limits>
#include <pollux/exec/aggregate.h>
#include <pollux/exec/aggregation_hook.h>
#include <pollux/functions/lib/aggregates/min_nax_aggregate_base.h>
#include <pollux/functions/lib/aggregates/simple_numeric_aggregate.h>
#include <pollux/functions/lib/aggregates/value_set.h>
#include <pollux/functions/prestosql/aggregates/aggregate_names.h>
#include <pollux/type/floating_point_util.h>

using namespace kumo::pollux::functions::aggregate;

namespace kumo::pollux::aggregate::prestosql {
    namespace {
        std::pair<vector_size_t *, vector_size_t *> rawOffsetAndSizes(
            ArrayVector &array_vector) {
            return {
                array_vector.offsets()->asMutable<vector_size_t>(),
                array_vector.sizes()->asMutable<vector_size_t>()
            };
        }

        /// @tparam V Type of value.
        /// @tparam Compare Type of comparator for T.
        template<typename T, typename Compare>
        struct MinMaxNAccumulator {
            int64_t n{0};
            using Allocator = std::conditional_t<
                std::is_same_v<int128_t, T>,
                AlignedStlAllocator<T, sizeof(int128_t)>,
                StlAllocator<T> >;
            std::vector<T, Allocator> heapValues;

            explicit MinMaxNAccumulator(HashStringAllocator *allocator)
                : heapValues{Allocator(allocator)} {
            }

            int64_t getN() const {
                return n;
            }

            size_t size() const {
                return heapValues.size();
            }

            void checkAndSetN(DecodedVector &decodedN, vector_size_t row) {
                // Skip null N.
                if (decodedN.is_null_at(row)) {
                    return;
                }

                const auto newN = decodedN.value_at<int64_t>(row);
                POLLUX_USER_CHECK_GT(
                    newN, 0, "second argument of max/min must be a positive integer");

                POLLUX_USER_CHECK_LE(
                    newN,
                    10'000,
                    "second argument of max/min must be less than or equal to 10000");

                if (n) {
                    POLLUX_USER_CHECK_EQ(
                        newN,
                        n,
                        "second argument of max/min must be a constant for all rows in a group");
                } else {
                    n = newN;
                }
            }

            void compareAndAdd(T value, Compare &comparator) {
                if (heapValues.size() < n) {
                    heapValues.push_back(value);
                    std::push_heap(heapValues.begin(), heapValues.end(), comparator);
                } else {
                    const auto &topValue = heapValues.front();
                    if (comparator(value, topValue)) {
                        std::pop_heap(heapValues.begin(), heapValues.end(), comparator);
                        heapValues.back() = value;
                        std::push_heap(heapValues.begin(), heapValues.end(), comparator);
                    }
                }
            }

            /// Copy all values from 'topValues' into 'values'. The heap remains
            /// unchanged after the call.
            void
            extractValues(VectorPtr &values, vector_size_t offset, Compare &comparator) {
                auto rawValues = values->as_flat_vector<T>()->mutableRawValues();
                std::sort_heap(heapValues.begin(), heapValues.end(), comparator);
                for (int64_t i = heapValues.size() - 1; i >= 0; --i) {
                    rawValues[offset + i] = heapValues[i];
                }
                std::make_heap(heapValues.begin(), heapValues.end(), comparator);
            }
        };

        /// @tparam Compare Type of comparator for T.
        template<typename Compare>
        struct MinMaxNAccumulator<StringView, Compare> {
            int64_t n{0};
            using Allocator = StlAllocator<StringView>;
            std::vector<StringView, Allocator> heapValues;
            ValueSet valueSet;

            explicit MinMaxNAccumulator(HashStringAllocator *allocator)
                : heapValues{Allocator(allocator)}, valueSet{allocator} {
            }

            int64_t getN() const {
                return n;
            }

            size_t size() const {
                return heapValues.size();
            }

            void checkAndSetN(DecodedVector &decodedN, vector_size_t row) {
                // Skip null N.
                if (decodedN.is_null_at(row)) {
                    return;
                }

                const auto newN = decodedN.value_at<int64_t>(row);
                POLLUX_USER_CHECK_GT(
                    newN, 0, "second argument of max/min must be a positive integer");

                POLLUX_USER_CHECK_LE(
                    newN,
                    10'000,
                    "second argument of max/min must be less than or equal to 10000");

                if (n) {
                    POLLUX_USER_CHECK_EQ(
                        newN,
                        n,
                        "second argument of max/min must be a constant for all rows in a group");
                } else {
                    n = newN;
                }
            }

            void compareAndAdd(StringView value, Compare &comparator) {
                if (heapValues.size() < n) {
                    heapValues.push_back(valueSet.write(value));
                    std::push_heap(heapValues.begin(), heapValues.end(), comparator);
                } else {
                    const auto &topValue = heapValues.front();
                    if (comparator(value, topValue)) {
                        std::pop_heap(heapValues.begin(), heapValues.end(), comparator);
                        valueSet.free(heapValues.back());
                        heapValues.back() = valueSet.write(value);
                        std::push_heap(heapValues.begin(), heapValues.end(), comparator);
                    }
                }
            }

            /// Copy all values from 'topValues' into 'values'. The heap remains
            /// unchanged after the call.
            void
            extractValues(VectorPtr &values, vector_size_t offset, Compare &comparator) {
                auto result = values->as_flat_vector<StringView>();
                std::sort_heap(heapValues.begin(), heapValues.end(), comparator);
                for (int64_t i = heapValues.size() - 1; i >= 0; --i) {
                    result->set(offset + i, heapValues[i]);
                }
                std::make_heap(heapValues.begin(), heapValues.end(), comparator);
            }
        };

        template<typename T, typename Compare>
        class MinMaxNAggregateBase : public exec::Aggregate {
        protected:
            explicit MinMaxNAggregateBase(const TypePtr &resultType)
                : exec::Aggregate(resultType) {
            }

            using AccumulatorType = MinMaxNAccumulator<T, Compare>;

            bool isFixedSize() const override {
                return false;
            }

            int32_t accumulatorFixedWidthSize() const override {
                return sizeof(AccumulatorType);
            }

            void initializeNewGroupsInternal(
                char **groups,
                melon::Range<const vector_size_t *> indices) override {
                exec::Aggregate::setAllNulls(groups, indices);
                for (const vector_size_t i: indices) {
                    auto group = groups[i];
                    new(group + offset_) AccumulatorType(allocator_);
                }
            }

            void addRawInput(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*unused*/) override {
                decodedValue_.decode(*args[0], rows);
                decodedN_.decode(*args[1], rows);

                rows.applyToSelected([&](vector_size_t i) {
                    if (decodedValue_.is_null_at(i) || decodedN_.is_null_at(i)) {
                        return;
                    }

                    auto *group = groups[i];

                    auto *accumulator = value(group);
                    accumulator->checkAndSetN(decodedN_, i);

                    auto tracker = trackRowSize(group);
                    addRawInput(group, i);
                });
            }

            void addIntermediateResults(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                auto results = decodeIntermediateResults(args[0], rows);

                rows.applyToSelected([&](vector_size_t i) {
                    if (!decodedIntermediates_.is_null_at(i)) {
                        addIntermediateResults(groups[i], i, results);
                    }
                });
            }

            void addSingleGroupRawInput(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*unused*/) override {
                decodedValue_.decode(*args[0], rows);

                auto *accumulator = value(group);
                validateN(args[1], rows, accumulator);

                auto tracker = trackRowSize(group);
                rows.applyToSelected([&](vector_size_t i) {
                    // Skip null value or N.
                    if (!decodedValue_.is_null_at(i) && !decodedN_.is_null_at(i)) {
                        addRawInput(group, i);
                    }
                });
            }

            void addSingleGroupIntermediateResults(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                auto results = decodeIntermediateResults(args[0], rows);

                rows.applyToSelected([&](vector_size_t i) {
                    if (!decodedIntermediates_.is_null_at(i)) {
                        addIntermediateResults(group, i, results);
                    }
                });
            }

            void extractValues(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                auto valuesArray = (*result)->as<ArrayVector>();
                valuesArray->resize(numGroups);

                const auto numValues =
                        countValuesAndSetResultNulls(groups, numGroups, *result);

                auto values = valuesArray->elements();
                values->resize(numValues);

                auto [rawOffsets, rawSizes] = rawOffsetAndSizes(*valuesArray);

                extractValues(groups, numGroups, rawOffsets, rawSizes, values, nullptr);
            }

            void extractAccumulators(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                auto row_vector = (*result)->as<RowVector>();
                row_vector->resize(numGroups);

                auto nVector = row_vector->childAt(0);
                nVector->resize(numGroups);

                auto valuesArray = row_vector->childAt(1)->as<ArrayVector>();
                valuesArray->resize(numGroups);

                const auto numValues =
                        countValuesAndSetResultNulls(groups, numGroups, *result);

                auto values = valuesArray->elements();
                values->resize(numValues);

                auto *rawNs = nVector->as<FlatVector<int64_t> >()->mutableRawValues();

                auto [rawOffsets, rawSizes] = rawOffsetAndSizes(*valuesArray);

                extractValues(groups, numGroups, rawOffsets, rawSizes, values, rawNs);
            }

            void destroyInternal(melon::Range<char **> groups) override {
                destroyAccumulators<AccumulatorType>(groups);
            }

        private:
            inline AccumulatorType *value(char *group) {
                return reinterpret_cast<AccumulatorType *>(group + Aggregate::offset_);
            }

            void extractValues(
                char **groups,
                int32_t numGroups,
                vector_size_t *rawOffsets,
                vector_size_t *rawSizes,
                VectorPtr &values,
                int64_t *rawNs) {
                vector_size_t offset = 0;
                for (auto i = 0; i < numGroups; ++i) {
                    auto *group = groups[i];

                    if (!isNull(group)) {
                        auto *accumulator = value(group);
                        const vector_size_t size = accumulator->size();

                        rawOffsets[i] = offset;
                        rawSizes[i] = size;

                        if (rawNs != nullptr) {
                            rawNs[i] = accumulator->n;
                        }
                        accumulator->extractValues(values, offset, comparator_);

                        offset += size;
                    }
                }
            }

            void addRawInput(char *group, vector_size_t index) {
                clearNull(group);

                auto *accumulator = value(group);

                accumulator->compareAndAdd(decodedValue_.value_at<T>(index), comparator_);
            }

            struct IntermediateResult {
                const ArrayVector *valueArray;
                const FlatVector<T> *flatValues;
            };

            void addIntermediateResults(
                char *group,
                vector_size_t index,
                IntermediateResult &result) {
                clearNull(group);

                auto *accumulator = value(group);

                const auto decodedIndex = decodedIntermediates_.index(index);

                accumulator->checkAndSetN(decodedN_, decodedIndex);

                const auto *valueArray = result.valueArray;
                const auto *values = result.flatValues;
                auto *rawValues = values->rawValues();

                const auto numValues = valueArray->sizeAt(decodedIndex);
                const auto valueOffset = valueArray->offsetAt(decodedIndex);

                auto tracker = trackRowSize(group);
                for (auto i = 0; i < numValues; ++i) {
                    const auto v = rawValues[valueOffset + i];
                    accumulator->compareAndAdd(v, comparator_);
                }
            }

            IntermediateResult decodeIntermediateResults(
                const VectorPtr &arg,
                const SelectivityVector &rows) {
                decodedIntermediates_.decode(*arg, rows);

                auto baseRowVector =
                        dynamic_cast<const RowVector *>(decodedIntermediates_.base());

                decodedN_.decode(*baseRowVector->childAt(0), rows);
                decodedValue_.decode(*baseRowVector->childAt(1), rows);

                IntermediateResult result{};
                result.valueArray = decodedValue_.base()->template as<ArrayVector>();
                result.flatValues =
                        result.valueArray->elements()->template as<FlatVector<T> >();

                return result;
            }

            /// Return total number of values in all accumulators of specified 'groups'.
            /// Set null flags in 'result'.
            vector_size_t countValuesAndSetResultNulls(
                char **groups,
                int32_t numGroups,
                VectorPtr &result) {
                vector_size_t numValues = 0;

                uint64_t *rawNulls = getRawNulls(result.get());

                for (auto i = 0; i < numGroups; ++i) {
                    auto *group = groups[i];
                    auto *accumulator = value(group);

                    if (isNull(group)) {
                        result->set_null(i, true);
                    } else {
                        clearNull(rawNulls, i);
                        numValues += accumulator->size();
                    }
                }

                return numValues;
            }

            void validateN(
                const VectorPtr &arg,
                const SelectivityVector &rows,
                AccumulatorType *accumulator) {
                decodedN_.decode(*arg, rows);
                if (decodedN_.isConstantMapping()) {
                    accumulator->checkAndSetN(decodedN_, rows.begin());
                } else {
                    rows.applyToSelected(
                        [&](auto row) { accumulator->checkAndSetN(decodedN_, row); });
                }
            }

            Compare comparator_;
            DecodedVector decodedValue_;
            DecodedVector decodedN_;
            DecodedVector decodedIntermediates_;
        };

        template<typename T>
        struct LessThanComparator : public std::less<T> {
        };

        template<>
        struct LessThanComparator<float>
                : public util::floating_point::NaNAwareLessThan<float> {
        };

        template<>
        struct LessThanComparator<double>
                : public util::floating_point::NaNAwareLessThan<double> {
        };

        template<typename T>
        struct GreaterThanComparator : public std::greater<T> {
        };

        template<>
        struct GreaterThanComparator<float>
                : public util::floating_point::NaNAwareGreaterThan<float> {
        };

        template<>
        struct GreaterThanComparator<double>
                : public util::floating_point::NaNAwareGreaterThan<double> {
        };

        template<typename T>
        class MinNAggregate : public MinMaxNAggregateBase<T, LessThanComparator<T> > {
        public:
            explicit MinNAggregate(const TypePtr &resultType)
                : MinMaxNAggregateBase<T, LessThanComparator<T> >(resultType) {
            }
        };

        template<typename T>
        class MaxNAggregate : public MinMaxNAggregateBase<T, GreaterThanComparator<T> > {
        public:
            explicit MaxNAggregate(const TypePtr &resultType)
                : MinMaxNAggregateBase<T, GreaterThanComparator<T> >(resultType) {
            }
        };

        template<template <typename T> typename AggregateN>
        exec::AggregateRegistrationResult registerMinMax(
            const std::string &name,
            bool withCompanionFunctions,
            bool overwrite,
            bool registerMin) {
            std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures;
            signatures.push_back(exec::AggregateFunctionSignatureBuilder()
                .orderableTypeVariable("T")
                .returnType("T")
                .intermediateType("T")
                .argumentType("T")
                .build());
            for (const auto &type:
                 {
                     "tinyint",
                     "integer",
                     "smallint",
                     "bigint",
                     "real",
                     "double",
                     "varchar"
                 }) {
                // T, bigint -> row(array(T), bigint) -> array(T)
                signatures.push_back(
                    exec::AggregateFunctionSignatureBuilder()
                    .returnType(fmt::format("array({})", type))
                    .intermediateType(fmt::format("row(bigint, array({}))", type))
                    .argumentType(type)
                    .argumentType("bigint")
                    .build());
            }

            // decimal(p,s), bigint -> row(array(decimal(p,s)), bigint) ->
            // array(decimal(p,s))
            signatures.push_back(
                exec::AggregateFunctionSignatureBuilder()
                .integerVariable("a_precision")
                .integerVariable("a_scale")
                .argumentType("DECIMAL(a_precision, a_scale)")
                .argumentType("bigint")
                .intermediateType("row(bigint, array(DECIMAL(a_precision, a_scale)))")
                .returnType("array(DECIMAL(a_precision, a_scale))")
                .build());

            return exec::registerAggregateFunction(
                name,
                std::move(signatures),
                [name, registerMin](
            core::AggregationNode::Step step,
            std::vector<TypePtr> argTypes,
            const TypePtr &resultType,
            const core::QueryConfig &config) -> std::unique_ptr<exec::Aggregate> {
                    const bool nAgg = !resultType->equivalent(*argTypes[0]);
                    if (nAgg) {
                        // We have either 2 arguments: T, bigint (partial aggregation)
                        // or one argument: row(bigint, array(T)) (intermediate or final
                        // aggregation). Extract T.
                        const auto &inputType = argTypes.size() == 2
                                                    ? argTypes[0]
                                                    : argTypes[0]->childAt(1)->childAt(0);

                        switch (inputType->kind()) {
                            case TypeKind::TINYINT:
                                return std::make_unique<AggregateN<int8_t> >(resultType);
                            case TypeKind::SMALLINT:
                                return std::make_unique<AggregateN<int16_t> >(resultType);
                            case TypeKind::INTEGER:
                                return std::make_unique<AggregateN<int32_t> >(resultType);
                            case TypeKind::BIGINT:
                                return std::make_unique<AggregateN<int64_t> >(resultType);
                            case TypeKind::REAL:
                                return std::make_unique<AggregateN<float> >(resultType);
                            case TypeKind::DOUBLE:
                                return std::make_unique<AggregateN<double> >(resultType);
                            case TypeKind::TIMESTAMP:
                                return std::make_unique<AggregateN<Timestamp> >(resultType);
                            case TypeKind::VARCHAR:
                                return std::make_unique<AggregateN<StringView> >(resultType);
                            case TypeKind::HUGEINT:
                                if (inputType->isLongDecimal()) {
                                    return std::make_unique<AggregateN<int128_t> >(resultType);
                                }
                                [[fallthrough]];
                            default:
                                POLLUX_UNREACHABLE(
                                    "Unknown input type for {} aggregation {}",
                                    name,
                                    inputType->kindName());
                        }
                    } else {
                        if (registerMin) {
                            auto factory = getMinFunctionFactory(
                                name,
                                CompareFlags::NullHandlingMode::kNullAsIndeterminate,
                                TimestampPrecision::kMilliseconds);
                            return factory(step, argTypes, resultType, config);
                        } else {
                            auto factory = getMaxFunctionFactory(
                                name,
                                CompareFlags::NullHandlingMode::kNullAsIndeterminate,
                                TimestampPrecision::kMilliseconds);
                            return factory(step, argTypes, resultType, config);
                        }
                    }
                },
                {false /*orderSensitive*/, false /*companionFunction*/},
                withCompanionFunctions,
                overwrite);
        }
    } // namespace

    void registerMinMaxAggregates(
        const std::string &prefix,
        bool withCompanionFunctions,
        bool overwrite) {
        registerMinMax<MinNAggregate>(
            prefix + kMin, withCompanionFunctions, overwrite, true);
        registerMinMax<MaxNAggregate>(
            prefix + kMax, withCompanionFunctions, overwrite, false);
    }
} // namespace kumo::pollux::aggregate::prestosql
